import paramiko  # 导入paramiko库，用于SSH和SFTP连接
import logging  # 导入logging库，用于记录日志
import json  # 导入json库，用于序列化和反序列化消息
from contextlib import contextmanager  # 导入contextmanager，用于创建上下文管理器
from typing import Optional, List, Iterator  # 导入类型注解
from kafka import KafkaConsumer, KafkaProducer  # 导入Kafka的生产者和消费者类

logger = logging.getLogger(__name__)  # 获取当前模块的logger实例


class KafkaClient:
    """
    Kafka生产者和消费者的封装工具，基于kafka-python库

    用法示例：
        client = KafkaClient(bootstrap_servers, username, password)
        with client.get_producer() as producer:
            producer.send('topic', {'key': 'value'})
        for msg in client.consume('topic', group_id='my_group'):
            print(msg)
    """

    def __init__(
            self,
            bootstrap_servers: List[str],  # Kafka服务器地址列表
            username: str,  # SASL认证用户名
            password: str,  # SASL认证密码
            security_protocol: str = 'SASL_PLAINTEXT',  # 安全协议
            sasl_mechanism: str = 'PLAIN',  # SASL机制
            api_version: tuple = (2, 2, 10)  # Kafka API版本
    ) -> None:
        self.bootstrap_servers = bootstrap_servers  # 保存Kafka服务器地址
        self.username = username  # 保存用户名
        self.password = password  # 保存密码
        self.security_protocol = security_protocol  # 保存安全协议
        self.sasl_mechanism = sasl_mechanism  # 保存SASL机制
        self.api_version = api_version  # 保存API版本
        self._producer: Optional[KafkaProducer] = None  # 延迟初始化生产者实例
        self._consumer: Optional[KafkaConsumer] = None  # 延迟初始化消费者实例

    def _common_config(self) -> dict:
        # 构建通用配置字典
        return {
            'bootstrap_servers': self.bootstrap_servers,  # Kafka服务器地址
            'security_protocol': self.security_protocol,  # 安全协议
            'sasl_mechanism': self.sasl_mechanism,  # SASL机制
            'sasl_plain_username': self.username,  # SASL用户名
            'sasl_plain_password': self.password,  # SASL密码
            'api_version': self.api_version,  # API版本
        }

    @contextmanager
    def get_producer(
            self,
            key_serializer=lambda v: json.dumps(v).encode('utf-8'),  # 键序列化函数
            value_serializer=lambda v: json.dumps(v).encode('utf-8')  # 值序列化函数
    ) -> KafkaProducer:
        """
        上下文管理器，获取并管理KafkaProducer生命周期
        """
        if not self._producer:
            config = self._common_config()  # 获取通用配置
            config.update({
                'key_serializer': key_serializer,  # 注入键序列化
                'value_serializer': value_serializer,  # 注入值序列化
            })
            self._producer = KafkaProducer(**config)  # 创建生产者实例
            logger.info('KafkaProducer created')  # 记录日志：生产者创建
        try:
            yield self._producer  # 返回生产者给调用方
        finally:
            self._producer.flush()  # 刷新缓冲区，确保消息发送
            self._producer.close()  # 关闭生产者连接
            logger.info('KafkaProducer closed')  # 记录日志：生产者关闭
            self._producer = None  # 重置生产者实例

    def consume(
            self,
            topic: str,  # 主题名称
            group_id: str,  # 消费者组ID
            auto_offset_reset: str = 'latest',  # 偏移量重置策略
            enable_auto_commit: bool = False,  # 是否自动提交偏移量
            key_deserializer=lambda m: json.loads(m.decode('utf-8')),  # 键反序列化
            value_deserializer=lambda m: json.loads(m.decode('utf-8')),  # 值反序列化
    ) -> Iterator[dict]:
        """
        消费消息生成器，迭代返回Kafka消息对象
        """
        config = self._common_config()  # 获取通用配置
        config.update({
            'group_id': group_id,  # 设置消费者组
            'auto_offset_reset': auto_offset_reset,  # 偏移量重置
            'enable_auto_commit': enable_auto_commit,  # 自动提交
            'key_deserializer': key_deserializer,  # 注入键反序列化
            'value_deserializer': value_deserializer,  # 注入值反序列化
        })
        consumer = KafkaConsumer(topic, **config)  # 创建消费者实例并订阅主题
        logger.info(f'KafkaConsumer subscribed to {topic} with group {group_id}')  # 记录订阅日志
        try:
            for message in consumer:  # 遍历消息队列
                yield message  # 产出消息给调用方
        except Exception as e:
            logger.error(f'Error consuming messages: {e}')  # 记录错误日志
            raise  # 抛出异常给上层处理
        finally:
            consumer.close()  # 关闭消费者连接
            logger.info('KafkaConsumer closed')  # 记录消费者关闭日志
